import torch.nn as nn
import torch
import matplotlib.pyplot as plt

# 定义自注意力模块类，继承自 nn.Module
class Self_Attention(nn.Module):
    """
    自注意力模块类，用于计算自注意力机制。

    参数:
    dim (int): 输入特征的维度。
    dk (int): 查询（query）和键（key）的维度。
    dv (int): 值（value）的维度。
    """
    def __init__(self, dim, dk, dv):
        # 调用父类的构造函数
        super(Self_Attention, self).__init__()
        # 缩放因子，用于缩放注意力分数
        self.scale = dk ** -0.5
        # 定义线性层，用于将输入转换为查询向量
        self.q = nn.Linear(dim, dk)
        # 定义线性层，用于将输入转换为键向量
        self.k = nn.Linear(dim, dk)
        # 定义线性层，用于将输入转换为值向量
        self.v = nn.Linear(dim, dv)

    def forward(self, x):
        """
        前向传播方法，计算自注意力机制。

        参数:
        x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, dim)。

        返回:
        torch.Tensor: 经过自注意力计算后的输出张量，形状为 (batch_size, seq_len, dv)。
        """
        # 通过线性层得到查询向量
        q = self.q(x)
        print('q=',q)
        # 通过线性层得到键向量
        k = self.k(x)
        print('k=',k)
        # 通过线性层得到值向量
        v = self.v(x)
        print('v=',v)
        # 计算查询和键的点积，并乘以缩放因子
        attn = (q @ k.transpose(-2, -1)) * self.scale
        # 对注意力分数进行 softmax 操作，得到注意力分布
        print('attn=',attn)
        attn = attn.softmax(dim=-1)
        print('attn=',attn)

        # 将注意力分布与值向量相乘，得到输出
        x = attn @ v
        print('x=',x)
        return x

# 创建自注意力模块的实例，输入维度为 2，查询和键的维度为 2，值的维度为 3
att = Self_Attention(dim=2, dk=2, dv=3)

# 生成一个形状为 (1, 4, 2) 的随机张量，用于作为自注意力模块的输入
x = torch.rand((1, 4, 2))
print('x=',x)
# 将随机张量输入自注意力模块，得到输出
output = att(x)


#@save
class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        # 定义 Dropout 层，防止过拟合
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的位置编码矩阵 P
        self.P = torch.zeros((1, max_len, num_hiddens))
        # 计算用于生成位置编码的 X
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 将正弦函数结果赋值给 P 的偶数位置
        self.P[:, :, 0::2] = torch.sin(X)
        # 将余弦函数结果赋值给 P 的奇数位置
        self.P[:, :, 1::2] = torch.cos(X)
 
    def forward(self, X):
        # 将输入 X 与对应长度的位置编码相加
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        # 对加上位置编码后的 X 应用 Dropout 操作并返回
        return self.dropout(X)